{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7d4fbb7b-3ab3-4cc1-b567-bd95dc455931",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting elasticsearch\n",
      "  Downloading elasticsearch-8.15.0-py3-none-any.whl.metadata (8.7 kB)\n",
      "Collecting elastic-transport<9,>=8.13 (from elasticsearch)\n",
      "  Downloading elastic_transport-8.15.0-py3-none-any.whl.metadata (3.6 kB)\n",
      "Requirement already satisfied: urllib3<3,>=1.26.2 in /Users/richardhuth/Development/langwatch-sass/langwatch/langwatch_nlp/.venv/lib/python3.11/site-packages (from elastic-transport<9,>=8.13->elasticsearch) (2.2.1)\n",
      "Requirement already satisfied: certifi in /Users/richardhuth/Development/langwatch-sass/langwatch/langwatch_nlp/.venv/lib/python3.11/site-packages (from elastic-transport<9,>=8.13->elasticsearch) (2024.2.2)\n",
      "Downloading elasticsearch-8.15.0-py3-none-any.whl (523 kB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m523.3/523.3 kB\u001b[0m \u001b[31m6.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading elastic_transport-8.15.0-py3-none-any.whl (64 kB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.4/64.4 kB\u001b[0m \u001b[31m7.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: elastic-transport, elasticsearch\n",
      "Successfully installed elastic-transport-8.15.0 elasticsearch-8.15.0\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.0\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a8bfd342",
   "metadata": {},
   "outputs": [],
   "source": [
    "project_id = \"KAXYxPR8MUgTcP8CF193y\"\n",
    "# project_id = \"EEisP6epvj-no_veGiHTQ\"\n",
    "# project_id = \"project_iCTt0LSMXYbv5jZNSdtEr\"\n",
    "# project_id = \"project_FPGv6-1y-Sr4nqP1EoEQW\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "68573900-12e1-4865-9a4f-5f9e05d46f5c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>trace_id</th>\n",
       "      <th>input</th>\n",
       "      <th>embeddings</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>trace_kbsOG-U46Br8jv9TaE6Zq</td>\n",
       "      <td>hi</td>\n",
       "      <td>[-0.0037493026, -0.019173345, 0.0121757155, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>trace_aWlIO_1Pq8eW3R17BoiUK</td>\n",
       "      <td>hii</td>\n",
       "      <td>[0.038733196, -0.037454136, 0.025052842, 0.017...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>trace_--palpjWmJs_23f1y4D1u</td>\n",
       "      <td>aaaa</td>\n",
       "      <td>[0.02986709, -0.016567731, -0.008210792, 0.013...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>trace_TaHab1z-GJXOvzVWXWu57</td>\n",
       "      <td>hi</td>\n",
       "      <td>[-0.0037493026, -0.019173345, 0.0121757155, 0....</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>trace_3y_a8PMk-5xrSh0a8wz0u</td>\n",
       "      <td>hi</td>\n",
       "      <td>[-0.0037493026, -0.019173345, 0.0121757155, 0....</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                      trace_id input  \\\n",
       "0  trace_kbsOG-U46Br8jv9TaE6Zq    hi   \n",
       "1  trace_aWlIO_1Pq8eW3R17BoiUK   hii   \n",
       "2  trace_--palpjWmJs_23f1y4D1u  aaaa   \n",
       "3  trace_TaHab1z-GJXOvzVWXWu57    hi   \n",
       "4  trace_3y_a8PMk-5xrSh0a8wz0u    hi   \n",
       "\n",
       "                                          embeddings  \n",
       "0  [-0.0037493026, -0.019173345, 0.0121757155, 0....  \n",
       "1  [0.038733196, -0.037454136, 0.025052842, 0.017...  \n",
       "2  [0.02986709, -0.016567731, -0.008210792, 0.013...  \n",
       "3  [-0.0037493026, -0.019173345, 0.0121757155, 0....  \n",
       "4  [-0.0037493026, -0.019173345, 0.0121757155, 0....  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "import os\n",
    "from elasticsearch import Elasticsearch\n",
    "import pandas as pd\n",
    "\n",
    "es = Elasticsearch(\n",
    "    os.environ[\"ELASTICSEARCH_NODE_URL\"], api_key=os.environ[\"ELASTICSEARCH_API_KEY\"]\n",
    ")\n",
    "\n",
    "response = es.search(\n",
    "    index=\"search-traces-alias\",\n",
    "    query={\n",
    "        \"bool\": {\n",
    "            \"must\": [\n",
    "                {\"term\": {\"project_id\": project_id}},\n",
    "                {\"exists\": {\"field\": \"input.embeddings.embeddings\"}},\n",
    "                {\n",
    "                    \"range\": {\n",
    "                        \"timestamps.inserted_at\": {\n",
    "                            \"gte\": \"now-1M\",\n",
    "                            \"lt\": \"now\",\n",
    "                        }\n",
    "                    }\n",
    "                },\n",
    "            ],\n",
    "        }\n",
    "    },\n",
    "    source=[\"trace_id\", \"input\"],  # type: ignore\n",
    "    sort=[{\"timestamps.inserted_at\": \"asc\"}, {\"trace_id\": \"asc\"}],\n",
    "    size=2000,\n",
    ")\n",
    "\n",
    "df = pd.DataFrame(\n",
    "    [\n",
    "        {\n",
    "            \"trace_id\": hit[\"_source\"][\"trace_id\"],\n",
    "            \"input\": hit[\"_source\"][\"input\"][\"value\"],\n",
    "            \"embeddings\": hit[\"_source\"][\"input\"][\"embeddings\"][\"embeddings\"],\n",
    "        }\n",
    "        for hit in response[\"hits\"][\"hits\"]\n",
    "    ]\n",
    ")\n",
    "df.to_csv(f\"./data/traces_for_topics_{project_id}.csv\", index=False)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d1f5543",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from typing import TypedDict\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "df = pd.read_csv(f\"notebooks/data/traces_for_topics_{project_id}.csv\")\n",
    "# convert embeddings column from string to list of floats\n",
    "df[\"embeddings\"] = df[\"embeddings\"].apply(\n",
    "    lambda x: list(map(float, x[1:-1].split(\", \")))\n",
    ")\n",
    "\n",
    "\n",
    "class Trace(TypedDict):\n",
    "    trace_id: str\n",
    "    input: str\n",
    "    embeddings: list[float]\n",
    "\n",
    "\n",
    "traces: list[Trace] = []\n",
    "for i, row in df.iterrows():\n",
    "    traces.append(\n",
    "        Trace(\n",
    "            trace_id=row[\"trace_id\"],\n",
    "            input=row[\"input\"],\n",
    "            embeddings=row[\"embeddings\"],\n",
    "        )\n",
    "    )\n",
    "traces"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "767c8f75",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.cluster.hierarchy import dendrogram, linkage\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "embeddings = np.array([np.array(x) for x in df[\"embeddings\"]])\n",
    "Z = linkage(embeddings, \"ward\")\n",
    "fig = plt.figure(figsize=(25, 10))\n",
    "dn = dendrogram(Z)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb2d6bbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.cluster.hierarchy import fcluster\n",
    "\n",
    "cophenetic_distances_for_topics = 4\n",
    "cophenetic_distances_for_subtopics = 2\n",
    "topic_ids = fcluster(Z, cophenetic_distances_for_topics, criterion=\"distance\")\n",
    "\n",
    "set(topic_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "375a4112",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "from scipy.spatial.distance import cdist\n",
    "\n",
    "\n",
    "def calculate_centroid_and_distance(samples) -> tuple[list[float], float]:\n",
    "    centroid = np.mean([np.array(item[\"embeddings\"]) for item in samples], axis=0)\n",
    "    distances = cdist(\n",
    "        [np.array(item[\"embeddings\"]) for item in samples],\n",
    "        np.array([centroid]),\n",
    "        \"cosine\",\n",
    "    ).flatten()\n",
    "    p95_distance = np.percentile(distances, 95).astype(float)\n",
    "\n",
    "    return centroid, p95_distance\n",
    "\n",
    "\n",
    "minimum_subtopics_per_topic = 1\n",
    "minimum_traces_per_topic = 5\n",
    "\n",
    "\n",
    "def build_hierarchy(\n",
    "    traces: list[Trace],\n",
    "    cophenetic_distance: int,\n",
    "    with_embeddings=True,\n",
    "    maximum_p95_distance: float = 1,\n",
    "    with_subtopics=True,\n",
    ") -> dict[str, dict[str, list[Trace]]]:\n",
    "    embeddings = [t[\"embeddings\"] for t in traces]\n",
    "    Z = linkage(embeddings, \"ward\")\n",
    "    topic_ids = fcluster(Z, cophenetic_distance, criterion=\"distance\")\n",
    "\n",
    "    # Dictionary to hold our two-level hierarchy\n",
    "    hierarchy = {}\n",
    "\n",
    "    # Iterate over each unique topic to create subtopics\n",
    "    for topic_id in set(topic_ids):\n",
    "        # Isolate samples that belong to the current topic\n",
    "        indices_in_topic = [i for i, t in enumerate(topic_ids) if t == topic_id]\n",
    "\n",
    "        traces_in_topic = [traces[i] for i in indices_in_topic]\n",
    "\n",
    "        # If there's less than the minimum number of traces, skip this topic\n",
    "        unique_inputs = set([t[\"input\"] for t in traces_in_topic])\n",
    "        if len(unique_inputs) < minimum_traces_per_topic:\n",
    "            continue\n",
    "\n",
    "        _, p95_distance = calculate_centroid_and_distance(traces_in_topic)\n",
    "        # Skip this topic if the p95 distance is too large\n",
    "        if p95_distance > maximum_p95_distance:\n",
    "            continue\n",
    "\n",
    "        if with_subtopics:\n",
    "            subtopics = build_hierarchy(\n",
    "                traces_in_topic,\n",
    "                cophenetic_distances_for_subtopics,\n",
    "                with_embeddings,\n",
    "                maximum_p95_distance,\n",
    "                False,\n",
    "            )\n",
    "\n",
    "            if len(subtopics.keys()) < minimum_subtopics_per_topic:\n",
    "                continue\n",
    "\n",
    "            hierarchy[f\"Topic {topic_id}\"] = subtopics\n",
    "        else:\n",
    "            hierarchy[f\"Subtopic {topic_id}\"] = [\n",
    "                Trace(\n",
    "                    trace_id=traces[i][\"trace_id\"],\n",
    "                    input=traces[i][\"input\"],\n",
    "                    embeddings=traces[i][\"embeddings\"] if with_embeddings else \"[embeddings]\",  # type: ignore\n",
    "                )\n",
    "                for i in indices_in_topic\n",
    "            ]\n",
    "\n",
    "    return hierarchy\n",
    "\n",
    "\n",
    "hierarchy = build_hierarchy(traces, cophenetic_distances_for_topics)\n",
    "pprint(build_hierarchy(traces, cophenetic_distances_for_topics, False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44524e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "topics_count = len(hierarchy.keys())\n",
    "subtopics_counts = [len(hierarchy[key].keys()) for key in hierarchy.keys()]\n",
    "samples_counts = [\n",
    "    len(samples) for subtopics in hierarchy.values() for samples in subtopics.values()\n",
    "]\n",
    "\n",
    "print(f\"Number of topics: {topics_count}\")\n",
    "print(f\"Number of subtopics: {subtopics_counts}\")\n",
    "print(\n",
    "    f\"Stats of samples:\",\n",
    "    \"Mean\",\n",
    "    np.mean(samples_counts),\n",
    "    \"Median\",\n",
    "    np.median(samples_counts),\n",
    "    \"Max\",\n",
    "    np.max(samples_counts),\n",
    "    \"Min\",\n",
    "    np.min(samples_counts),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "082f66f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "from openai import OpenAI\n",
    "import os\n",
    "import json\n",
    "\n",
    "openai = OpenAI(api_key=os.environ[\"OPENAI_API_KEY\"])\n",
    "\n",
    "\n",
    "def generate_topic_names(\n",
    "    topic_examples: list[list[str]], existing: Optional[list[str]] = None\n",
    ") -> list[str]:\n",
    "    example_count = sum([len(examples) for examples in topic_examples])\n",
    "    print(\n",
    "        f\"Generating names for {len(topic_examples)} topics with {example_count} examples total.\"\n",
    "    )\n",
    "    topic_examples_str = \"\\n\\n\\n\".join(\n",
    "        [\n",
    "            f\"# Topic {index} Samples\\n\\n\" + \"\\n\".join(samples)\n",
    "            for index, samples in enumerate(topic_examples)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    existing_str = \"\\n\".join(existing) if existing else \"\"\n",
    "    existing_message = (\n",
    "        f\"\\n\\nThose are the topics that already exist, avoid using any names that may overlap \\\n",
    "        in meaning with them, think how the current sample data is unique and different from those instead:\\n\\n\\\n",
    "            {existing_str}\"\n",
    "        if existing_str\n",
    "        else \"\"\n",
    "    )\n",
    "\n",
    "    response = openai.chat.completions.create(\n",
    "        model=\"gpt-4-1106-preview\",\n",
    "        temperature=0.1,\n",
    "        messages=[\n",
    "            {\n",
    "                \"role\": \"system\",\n",
    "                \"content\": f\"You are a highly knowledgeable assistant tasked with taxonomy for naming topics \\\n",
    "                    based on a list of examples. Provide a single, descriptive name for each topic. \\\n",
    "                    Topic names should not be similar to each other, as the data is already organized, \\\n",
    "                    the disambiguation between two similar topics should be clear from the name alone.\\\n",
    "                        {existing_message}\",\n",
    "            },\n",
    "            {\"role\": \"user\", \"content\": f\"{topic_examples_str}\"},\n",
    "        ],\n",
    "        tools=[\n",
    "            {\n",
    "                \"type\": \"function\",\n",
    "                \"function\": {\n",
    "                    \"name\": \"topicNames\",\n",
    "                    \"parameters\": {\n",
    "                        \"type\": \"object\",\n",
    "                        \"properties\": dict(\n",
    "                            [\n",
    "                                (f\"topic_{index}\", {\"type\": \"string\"})\n",
    "                                for index in range(len(topic_examples_str))\n",
    "                            ]\n",
    "                        ),\n",
    "                    },\n",
    "                    \"description\": \"use this function to name all the topics based on the examples provided\",\n",
    "                },\n",
    "            }\n",
    "        ],\n",
    "        tool_choice={\"type\": \"function\", \"function\": {\"name\": \"topicNames\"}},\n",
    "    )\n",
    "\n",
    "    return list(json.loads(response.choices[0].message.tool_calls[0].function.arguments).values())  # type: ignore\n",
    "\n",
    "\n",
    "topic_names = generate_topic_names(\n",
    "    [\n",
    "        [\"example1\", \"example2\"],\n",
    "        [\"foo\", \"bar\"],\n",
    "    ]\n",
    ")\n",
    "\n",
    "print(topic_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f771718",
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import random\n",
    "from typing import Iterable, TypeVar, Optional\n",
    "\n",
    "T = TypeVar(\"T\")\n",
    "\n",
    "\n",
    "def shuffled(x: Iterable[T]) -> list[T]:\n",
    "    return sorted(x, key=lambda _: random())\n",
    "\n",
    "\n",
    "def get_subtopic_samples(samples: list[Trace], n=5):\n",
    "    unique_values = list(set([item[\"input\"] for item in samples]))\n",
    "    return [item[0:140] for item in shuffled(unique_values)[0:n]]\n",
    "\n",
    "\n",
    "def generate_topic_and_subtopic_names(\n",
    "    hierarchy: dict[str, dict[str, list[Trace]]],\n",
    "    existing: Optional[list[str]] = None,\n",
    "    skip_topic_names: bool = False,\n",
    "):\n",
    "    topic_examples = [\n",
    "        shuffled(\n",
    "            [\n",
    "                item\n",
    "                for samples in subtopics.values()\n",
    "                for item in get_subtopic_samples(samples, n=5)\n",
    "            ]\n",
    "        )[0:30]\n",
    "        for subtopics in hierarchy.values()\n",
    "    ]\n",
    "\n",
    "    topic_names = (\n",
    "        list(hierarchy.keys())\n",
    "        if skip_topic_names\n",
    "        else generate_topic_names(topic_examples, existing=existing)\n",
    "    )\n",
    "\n",
    "    subtopic_names = [\n",
    "        generate_topic_names(\n",
    "            [get_subtopic_samples(samples, n=10) for samples in subtopics.values()],\n",
    "            existing=existing,\n",
    "        )\n",
    "        for subtopics in hierarchy.values()\n",
    "    ]\n",
    "\n",
    "    return topic_names, subtopic_names\n",
    "\n",
    "\n",
    "topic_names, subtopic_names = generate_topic_and_subtopic_names(hierarchy)\n",
    "\n",
    "topic_names, subtopic_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6e244ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional, TypedDict\n",
    "from scipy.spatial.distance import cdist\n",
    "import nanoid\n",
    "\n",
    "\n",
    "class Topic(TypedDict):\n",
    "    id: str\n",
    "    name: str\n",
    "    centroid: list[float]\n",
    "    p95_distance: float\n",
    "\n",
    "\n",
    "class Subtopic(Topic):\n",
    "    parent_id: str\n",
    "\n",
    "\n",
    "class TraceTopicMap(TypedDict):\n",
    "    trace_id: str\n",
    "    topic_id: Optional[str]\n",
    "    subtopic_id: Optional[str]\n",
    "\n",
    "\n",
    "def build_response(\n",
    "    hierarchy: dict[str, dict[str, list[Trace]]],\n",
    "    topic_names: list[str],\n",
    "    subtopic_names: list[list[str]],\n",
    ") -> tuple[list[Topic], list[Subtopic], list[TraceTopicMap]]:\n",
    "    topics: list[Topic] = []\n",
    "    subtopics: list[Subtopic] = []\n",
    "    traces: list[TraceTopicMap] = []\n",
    "\n",
    "    for topic_idx, topic in enumerate(hierarchy.values()):\n",
    "        topic_id = None\n",
    "\n",
    "        topic_samples = [\n",
    "            item for subtopic_samples in topic.values() for item in subtopic_samples\n",
    "        ]\n",
    "        unique_values = list(set([item[\"input\"] for item in topic_samples]))\n",
    "        if (\n",
    "            len(topic.values()) >= minimum_subtopics_per_topic\n",
    "            and len(unique_values) >= minimum_traces_per_topic\n",
    "        ):\n",
    "            topic_id = f\"topic_{nanoid.generate()}\"\n",
    "            topic_name = topic_names[topic_idx]\n",
    "            topic_centroid, topic_p95_distance = calculate_centroid_and_distance(\n",
    "                topic_samples\n",
    "            )\n",
    "\n",
    "            topics.append(\n",
    "                Topic(\n",
    "                    id=topic_id,\n",
    "                    name=topic_name,\n",
    "                    centroid=topic_centroid,\n",
    "                    p95_distance=topic_p95_distance,\n",
    "                )\n",
    "            )\n",
    "\n",
    "        for subtopic_idx, subtopic in enumerate(topic.values()):\n",
    "            subtopic_id = None\n",
    "\n",
    "            unique_values = list(set([item[\"input\"] for item in subtopic]))\n",
    "            if topic_id and len(unique_values) >= minimum_traces_per_topic:\n",
    "                subtopic_id = f\"subtopic_{nanoid.generate()}\"\n",
    "                subtopic_name = subtopic_names[topic_idx][subtopic_idx]\n",
    "                subtopic_centroid, subtopic_p95_distance = (\n",
    "                    calculate_centroid_and_distance(subtopic)\n",
    "                )\n",
    "\n",
    "                subtopics.append(\n",
    "                    Subtopic(\n",
    "                        id=subtopic_id,\n",
    "                        name=subtopic_name,\n",
    "                        centroid=subtopic_centroid,\n",
    "                        p95_distance=subtopic_p95_distance,\n",
    "                        parent_id=topic_id,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "            for trace in subtopic:\n",
    "                traces.append(\n",
    "                    TraceTopicMap(\n",
    "                        trace_id=trace[\"trace_id\"],\n",
    "                        topic_id=topic_id,\n",
    "                        subtopic_id=subtopic_id,\n",
    "                    )\n",
    "                )\n",
    "\n",
    "    return topics, subtopics, traces\n",
    "\n",
    "\n",
    "topics, subtopics, traces_ = build_response(hierarchy, topic_names, subtopic_names)\n",
    "\n",
    "(topics, subtopics, traces_)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "772fad22",
   "metadata": {},
   "source": [
    "# Incremental Clustering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea17af3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "U = TypeVar(\"U\", Topic, Subtopic)\n",
    "\n",
    "\n",
    "def get_matching_topic(trace: Trace, topics: list[U]) -> Optional[U]:\n",
    "    trace_embeddings = np.array(trace[\"embeddings\"])\n",
    "    centroid_distances = cdist(\n",
    "        [trace_embeddings],\n",
    "        np.array([t[\"centroid\"] for t in topics]),\n",
    "        \"cosine\",\n",
    "    ).flatten()\n",
    "\n",
    "    # Find the closest topic to assign the trace\n",
    "    sorted_topics = sorted(zip(topics, centroid_distances), key=lambda x: x[1])\n",
    "\n",
    "    if sorted_topics[0][1] <= sorted_topics[0][0][\"p95_distance\"] * 1.1:\n",
    "        return sorted_topics[0][0]\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "\n",
    "def assign_trace_to_topic(\n",
    "    trace: Trace, topics: list[Topic], subtopics: list[Subtopic]\n",
    ") -> TraceTopicMap:\n",
    "    topic_id = None\n",
    "    subtopic_id = None\n",
    "\n",
    "    matching_topic = get_matching_topic(trace, topics)\n",
    "    if matching_topic:\n",
    "        topic_id = matching_topic[\"id\"]\n",
    "        subtopics_ = [s for s in subtopics if s[\"parent_id\"] == topic_id]\n",
    "\n",
    "        matching_subtopic = (\n",
    "            get_matching_topic(trace, subtopics_) if len(subtopics_) > 0 else None\n",
    "        )\n",
    "\n",
    "        if matching_subtopic:\n",
    "            subtopic_id = matching_subtopic[\"id\"]\n",
    "\n",
    "    return TraceTopicMap(\n",
    "        trace_id=trace[\"trace_id\"],\n",
    "        topic_id=topic_id,\n",
    "        subtopic_id=subtopic_id,\n",
    "    )\n",
    "\n",
    "\n",
    "topic_matches = 0\n",
    "topic_mismatches = 0\n",
    "topic_no_matches = 0\n",
    "subtopic_matches = 0\n",
    "subtopic_mismatches = 0\n",
    "subtopic_no_matches = 0\n",
    "\n",
    "for new_trace_to_add in traces:\n",
    "    assigned_traces = [\n",
    "        t for t in traces_ if t[\"trace_id\"] == new_trace_to_add[\"trace_id\"]\n",
    "    ]\n",
    "    if not len(assigned_traces):\n",
    "        continue\n",
    "\n",
    "    assigned_trace = assigned_traces[0]\n",
    "\n",
    "    current_topic = [t for t in topics if t[\"id\"] == assigned_trace[\"topic_id\"]][0]\n",
    "    current_subtopic = [\n",
    "        t for t in subtopics if t[\"id\"] == assigned_trace[\"subtopic_id\"]\n",
    "    ][0]\n",
    "\n",
    "    trace_topic_map = assign_trace_to_topic(new_trace_to_add, topics, subtopics)\n",
    "\n",
    "    if trace_topic_map[\"topic_id\"] == assigned_trace[\"topic_id\"]:\n",
    "        topic_matches += 1\n",
    "    elif trace_topic_map[\"topic_id\"] is None:\n",
    "        topic_no_matches += 1\n",
    "    else:\n",
    "        topic_mismatches += 1\n",
    "\n",
    "    if trace_topic_map[\"subtopic_id\"] == assigned_trace[\"subtopic_id\"]:\n",
    "        subtopic_matches += 1\n",
    "    elif trace_topic_map[\"subtopic_id\"] is None:\n",
    "        subtopic_no_matches += 1\n",
    "    else:\n",
    "        subtopic_mismatches += 1\n",
    "\n",
    "topics_precision = topic_matches / (topic_matches + topic_mismatches)\n",
    "topics_recall = topic_matches / (topic_matches + topic_mismatches + topic_no_matches)\n",
    "topics_f1 = 2 * (topics_precision * topics_recall) / (topics_precision + topics_recall)\n",
    "print(\n",
    "    f\"Topics: Precision: {topics_precision:.3f}, Recall: {topics_recall:.3f}, F1: {topics_f1:.3f}\"\n",
    ")\n",
    "\n",
    "subtopics_precision = subtopic_matches / (subtopic_matches + subtopic_mismatches)\n",
    "subtopics_recall = subtopic_matches / (\n",
    "    subtopic_matches + subtopic_mismatches + subtopic_no_matches\n",
    ")\n",
    "subtopics_f1 = (\n",
    "    2\n",
    "    * (subtopics_precision * subtopics_recall)\n",
    "    / (subtopics_precision + subtopics_recall)\n",
    ")\n",
    "print(\n",
    "    f\"Subtopics: Precision: {subtopics_precision:.3f}, Recall: {subtopics_recall:.3f}, F1: {subtopics_f1:.3f}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc3e3e30",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.cluster.hierarchy import dendrogram, linkage\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "new_traces_to_assign = [\n",
    "    trace\n",
    "    for trace in traces\n",
    "    if not assign_trace_to_topic(trace, topics, subtopics)[\"topic_id\"]\n",
    "]\n",
    "\n",
    "\n",
    "def maybe_create_new_topics(\n",
    "    traces: list[Trace], topics: list[U], cophenetic_distances: int, with_subtopics=True\n",
    ") -> tuple[list[Topic], list[Subtopic], list[TraceTopicMap]]:\n",
    "    average_p95_distance = np.mean([t[\"p95_distance\"] for t in topics]).astype(float)\n",
    "    print(\"average_p95_distance\", average_p95_distance)\n",
    "    print(\"number of new traces to assign\", len(traces))\n",
    "    print([x[\"input\"] for x in traces])\n",
    "    centroid = np.mean([np.array(x[\"embeddings\"]) for x in traces], axis=0)\n",
    "    distances = cdist(\n",
    "        [np.array(x[\"embeddings\"]) for x in traces],\n",
    "        np.array([centroid]),\n",
    "        \"cosine\",\n",
    "    ).flatten()\n",
    "    print(\"new traces p95 distance\", np.percentile(distances, 95).astype(float))\n",
    "\n",
    "    embeddings = np.array([np.array(x[\"embeddings\"]) for x in traces])\n",
    "    Z = linkage(embeddings, \"ward\")\n",
    "    fig = plt.figure(figsize=(25, 10))\n",
    "    dn = dendrogram(Z)\n",
    "    plt.show()\n",
    "\n",
    "    new_hierarchy = build_hierarchy(\n",
    "        traces,\n",
    "        cophenetic_distances,\n",
    "        True,\n",
    "        maximum_p95_distance=average_p95_distance,\n",
    "        with_subtopics=with_subtopics,\n",
    "    )\n",
    "\n",
    "    if len(new_hierarchy.keys()) == 0:\n",
    "        return [], [], []\n",
    "\n",
    "    existing = [t[\"name\"] for t in topics]\n",
    "    if with_subtopics:\n",
    "        topic_names, subtopic_names = generate_topic_and_subtopic_names(new_hierarchy, existing=existing)\n",
    "    else:\n",
    "        new_hierarchy: dict[str, dict[str, list[Trace]]] = {\"New Sub Topics\": new_hierarchy}  # type: ignore\n",
    "        topic_names, subtopic_names = generate_topic_and_subtopic_names(new_hierarchy, existing=existing, skip_topic_names=True)\n",
    "    return build_response(new_hierarchy, topic_names, subtopic_names)\n",
    "\n",
    "\n",
    "new_topics, new_subtopics, new_traces_to_assign = maybe_create_new_topics(\n",
    "    new_traces_to_assign, topics, cophenetic_distances_for_topics\n",
    ")\n",
    "\n",
    "(new_topics, new_subtopics, new_traces_to_assign)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "574669a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "new_traces_to_assign_to_subtopics_map: dict[str, list[Trace]] = {}\n",
    "for trace in traces:\n",
    "    trace_topic_map = assign_trace_to_topic(trace, topics, subtopics)\n",
    "    topic_id = trace_topic_map[\"topic_id\"]\n",
    "    if topic_id and not trace_topic_map[\"subtopic_id\"]:\n",
    "        if topic_id not in new_traces_to_assign_to_subtopics_map:\n",
    "            new_traces_to_assign_to_subtopics_map[topic_id] = []\n",
    "        new_traces_to_assign_to_subtopics_map[topic_id].append(trace)\n",
    "\n",
    "for topic_id, traces_ in new_traces_to_assign_to_subtopics_map.items():\n",
    "    subtopics_ = [t for t in subtopics if t[\"parent_id\"] == topic_id]\n",
    "    _, new_subtopics_, new_traces_to_assign_ = maybe_create_new_topics(\n",
    "        traces_,\n",
    "        subtopics_,\n",
    "        cophenetic_distances_for_subtopics,\n",
    "        with_subtopics=False,\n",
    "    )\n",
    "    new_subtopics__ = [\n",
    "        Subtopic(\n",
    "            id=s[\"id\"],\n",
    "            name=s[\"name\"],\n",
    "            centroid=s[\"centroid\"],\n",
    "            p95_distance=s[\"p95_distance\"],\n",
    "            parent_id=topic_id,\n",
    "        )\n",
    "        for s in new_subtopics_\n",
    "    ]\n",
    "\n",
    "    new_subtopics += new_subtopics__\n",
    "    new_traces_to_assign += new_traces_to_assign_\n",
    "\n",
    "(new_topics, new_subtopics, new_traces_to_assign)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "164dcabc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
